import SPTAG
import numpy as np
import gc

# 读取fbin文件
def read_fbin(file_path):
    with open(file_path, 'rb') as f:
        num_vectors = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取向量数量
        dim = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取维度
        data = np.fromfile(f, dtype=np.float32).reshape(num_vectors, dim)  # 读取所有向量并重塑为矩阵
    return data

# 读取ground truth文件
def read_gt(file_path):
    with open(file_path, 'rb') as f:
        num_queries = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取查询数量
        top_k = np.fromfile(f, dtype=np.int32, count=1)[0]  # 读取top-k值
        gt_data = np.fromfile(f, dtype=np.int32).reshape(num_queries, top_k)  # 读取ground truth并重塑为矩阵
    return gt_data

# 评估检索结果
def evaluate_results(results, gt_data):
    correct = 0
    total = len(results) * len(results[0])
    
    for i, result in enumerate(results):
        correct += len(set(result).intersection(set(gt_data[i])))

    accuracy = correct / total
    return accuracy

# 主测试程序
def main(query_fbin, gt_file, index_file, k=10):
    # 加载查询向量
    print("Loading query file")
    queries = read_fbin(query_fbin)
    print("Loading query file done")
    
    # 加载ground truth
    print("Loading gt file")
    gt_data = read_gt(gt_file)
    print("Loading gt file done")
    
    # 加载索引
    print("Loading index file")
    index = SPTAG.AnnIndex.Load(index_file)
    print("Loading index file done")

    # 执行查询
    print("Start query")
    results = []
    for i, query in enumerate(queries):
        # print("Querying",i,"th query")
        result = index.Search(query, k)
        results.append(result[0])  # 提取检索结果中的ID
    print("query done")

    # 评估检索结果
    accuracy = evaluate_results(results, gt_data)
    print(f"Top-{k} accuracy: {accuracy:.4f}")

# 调用主程序
query_fbin = '/home/gary/Code/DiskANN/build/data/sift/sift_query.fbin'
gt_file = '/home/gary/Code/DiskANN/build/data/sift/sift_groundtruth.fbin'
index_file = 'sift1m_sptag_index'

main(query_fbin, gt_file, index_file, k=10)